package redis.limiter;

/**
 * @author zyl
 * @description
 * @since 2025/8/22
 */
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;

import java.util.Collections;

public class RedisTokenBucketRateLimiter {
    private final RedisTemplate<String, Object> redisTemplate;
    /**
     * Redis中存储令牌桶状态的键名，用于区分不同的限流维度
     */
    private final String key;
    /**
     * 桶的最大容量，决定突发流量的上限
     */
    private final int capacity;
    /**
     * 令牌补充速率（令牌数/秒），决定平均请求速率
     */
    private final double refillRate;

    public RedisTokenBucketRateLimiter(RedisTemplate<String, Object> redisTemplate,
                                       String key, int capacity, double refillRate) {
        this.redisTemplate = redisTemplate;
        this.key = key;
        this.capacity = capacity;
        this.refillRate = refillRate;
    }

    private static final String LUA_SCRIPT =
            "local key = KEYS[1]\n" +
                    "local now = tonumber(ARGV[1])\n" +
                    "local capacity = tonumber(ARGV[2])\n" +
                    "local refillRate = tonumber(ARGV[3])\n" +
                    "local requested = tonumber(ARGV[4])\n" +
                    "\n" +
                    "local bucket = redis.call('hmget', key, 'tokens', 'lastRefillTime')\n" +
                    "local tokens = tonumber(bucket[1])\n" +
                    "local lastRefillTime = tonumber(bucket[2])\n" +
                    "\n" +
                    "if tokens == nil or lastRefillTime == nil then\n" +
                    "    tokens = capacity\n" +
                    "    lastRefillTime = now\n" +
                    "end\n" +
                    "\n" +
                    "local timePassed = now - lastRefillTime\n" +
                    "if timePassed > 0 then\n" +
                    "    local newTokens = timePassed * refillRate\n" +
                    "    tokens = math.min(capacity, tokens + newTokens)\n" +
                    "    lastRefillTime = now\n" +
                    "end\n" +
                    "\n" +
                    "if tokens >= requested then\n" +
                    "    tokens = tokens - requested\n" +
                    "    redis.call('hmset', key, 'tokens', tokens, 'lastRefillTime', lastRefillTime)\n" +
                    "    redis.call('expire', key, math.ceil(capacity / refillRate) * 2)\n" +
                    "    return 1\n" +
                    "else\n" +
                    "    redis.call('hmset', key, 'tokens', tokens, 'lastRefillTime', lastRefillTime)\n" +
                    "    redis.call('expire', key, math.ceil(capacity / refillRate) * 2)\n" +
                    "    return 0\n" +
                    "end";

    public boolean tryAcquire(int tokens) {
        RedisScript<Long> script = new DefaultRedisScript<>(LUA_SCRIPT, Long.class);
        long now = System.currentTimeMillis() / 1000;
        Long result = redisTemplate.execute(script,
                Collections.singletonList(key),
                now, capacity, refillRate, tokens);
        return result != null && result == 1;
    }

    public boolean tryAcquire() {
        return tryAcquire(1);
    }
}